TypeScript Discriminated Unions: Building Type-Safe State Machines

In the realm of software development, managing application state effectively is crucial. State machines provide a powerful abstraction for modeling complex stateful systems, ensuring predictable behavior and simplifying reasoning about the system's logic. TypeScript, with its robust type system, offers a fantastic mechanism for building type-safe state machines using discriminated unions (also known as tagged unions or algebraic data types).

What are Discriminated Unions?

A discriminated union is a type that represents a value that can be one of several different types. Each of these types, known as members of the union, shares a common, distinct property called the discriminant or tag. This discriminant allows TypeScript to precisely determine which member of the union is currently active, enabling powerful type checking and auto-completion.

Think of it like a traffic light. It can be in one of three states: Red, Yellow, or Green. The 'color' property acts as the discriminant, telling us exactly which state the light is in.

Why Use Discriminated Unions for State Machines?

Discriminated unions bring several key benefits when building state machines in TypeScript:

• Type Safety: The compiler can verify that all possible states and transitions are handled correctly, preventing runtime errors related to unexpected state transitions. This is especially useful in large, complex applications.
• Exhaustiveness Checking: TypeScript can ensure that your code handles all possible states of the state machine, alerting you at compile time if a state is missed in a conditional statement or switch case. This helps to prevent unexpected behavior and makes your code more robust.
• Improved Readability: Discriminated unions clearly define the possible states of the system, making the code easier to understand and maintain. The explicit representation of states enhances code clarity.
• Enhanced Code Completion: TypeScript's intellisense provides intelligent code completion suggestions based on the current state, reducing the likelihood of errors and speeding up development.

Defining a State Machine with Discriminated Unions

Let's illustrate how to define a state machine using discriminated unions with a practical example: an order processing system. An order can be in the following states: Pending, Processing, Shipped, and Delivered.

Step 1: Define the State Types

First, we define the individual types for each state. Each type will have a `type` property acting as the discriminant, along with any state-specific data.

interface Pending {
  type: "pending";
  orderId: string;
  customerName: string;
  items: string[];
}

interface Processing {
  type: "processing";
  orderId: string;
  assignedAgent: string;
}

interface Shipped {
  type: "shipped";
  orderId: string;
  trackingNumber: string;
}

interface Delivered {
  type: "delivered";
  orderId: string;
  deliveryDate: Date;
}

Step 2: Create the Discriminated Union Type

Next, we create the discriminated union by combining these individual types using the `|` (union) operator.

type OrderState = Pending | Processing | Shipped | Delivered;

Now, `OrderState` represents a value that can be either `Pending`, `Processing`, `Shipped`, or `Delivered`. The `type` property within each state acts as the discriminant, allowing TypeScript to differentiate between them.

Handling State Transitions

Now that we have defined our state machine, we need a mechanism to transition between states. Let's create a `processOrder` function that takes the current state and an action as input and returns the new state.

interface Action {
  type: string;
  payload?: any;
}

function processOrder(state: OrderState, action: Action): OrderState {
  switch (state.type) {
    case "pending":
      if (action.type === "startProcessing") {
        return {
          type: "processing",
          orderId: state.orderId,
          assignedAgent: action.payload.agentId,
        };
      }
      return state; // No state change

    case "processing":
      if (action.type === "shipOrder") {
        return {
          type: "shipped",
          orderId: state.orderId,
          trackingNumber: action.payload.trackingNumber,
        };
      }
      return state; // No state change

    case "shipped":
      if (action.type === "deliverOrder") {
        return {
          type: "delivered",
          orderId: state.orderId,
          deliveryDate: new Date(),
        };
      }
      return state; // No state change

    case "delivered":
      // Order is already delivered, no further actions
      return state;

    default:
      // This should never happen due to exhaustiveness checking
      return state; // Or throw an error
  }
}

Explanation

• The `processOrder` function takes the current `OrderState` and an `Action` as input.
• It uses a `switch` statement to determine the current state based on the `state.type` discriminant.
• Inside each `case`, it checks the `action.type` to determine if a valid transition is triggered.
• If a valid transition is found, it returns a new state object with the appropriate `type` and data.
• If no valid transition is found, it returns the current state (or throws an error, depending on the desired behavior).
• The `default` case is included for completeness and should ideally never be reached due to TypeScript's exhaustiveness checking.

Leveraging Exhaustiveness Checking

TypeScript's exhaustiveness checking is a powerful feature that ensures you handle all possible states in your state machine. If you add a new state to the `OrderState` union but forget to update the `processOrder` function, TypeScript will flag an error.

To enable exhaustiveness checking, you can use the `never` type. Inside the `default` case of your switch statement, assign the state to a variable of type `never`.

function processOrder(state: OrderState, action: Action): OrderState {
  switch (state.type) {
    // ... (previous cases) ...

    default:
      const _exhaustiveCheck: never = state;
      return _exhaustiveCheck; // Or throw an error
  }
}

If the `switch` statement handles all possible `OrderState` values, the `_exhaustiveCheck` variable will be of type `never` and the code will compile. However, if you add a new state to the `OrderState` union and forget to handle it in the `switch` statement, the `_exhaustiveCheck` variable will be of a different type, and TypeScript will throw a compile-time error, alerting you to the missing case.

Practical Examples and Applications

Discriminated unions are applicable in a wide range of scenarios beyond simple order processing systems:

• UI State Management: Modeling the state of a UI component (e.g., loading, success, error).
• Network Request Handling: Representing the different stages of a network request (e.g., initial, in progress, success, failure).
• Form Validation: Tracking the validity of form fields and the overall form state.
• Game Development: Defining the different states of a game character or object.
• Authentication Flows: Managing user authentication states (e.g., logged in, logged out, pending verification).

Example: UI State Management

Let's consider a simple example of managing the state of a UI component that fetches data from an API. We can define the following states:

interface Initial {
  type: "initial";
}

interface Loading {
  type: "loading";
}

interface Success {
  type: "success";
  data: T;
}

interface Error {
  type: "error";
  message: string;
}

type UIState = Initial | Loading | Success | Error;

function renderUI(state: UIState): React.ReactNode {
  switch (state.type) {
    case "initial":
      return 
Click the button to load data.
;
    case "loading":
      return 
Loading...
;
    case "success":
      return 
{JSON.stringify(state.data, null, 2)}
;
    case "error":
      return 
Error: {state.message}
;
    default:
      const _exhaustiveCheck: never = state;
      return _exhaustiveCheck;
  }
}

This example demonstrates how discriminated unions can be used to effectively manage the different states of a UI component, ensuring that the UI is rendered correctly based on the current state. The `renderUI` function handles each state appropriately, providing a clear and type-safe way to manage the UI.

Best Practices for Using Discriminated Unions

To effectively utilize discriminated unions in your TypeScript projects, consider the following best practices:

• Choose Meaningful Discriminant Names: Select discriminant names that clearly indicate the purpose of the property (e.g., `type`, `state`, `status`).
• Keep State Data Minimal: Each state should only contain the data that is relevant to that specific state. Avoid storing unnecessary data in states.
• Use Exhaustiveness Checking: Always enable exhaustiveness checking to ensure that you handle all possible states.
• Consider Using a State Management Library: For complex state machines, consider using a dedicated state management library like XState, which provides advanced features such as state charts, hierarchical states, and parallel states. However, for simpler scenarios, discriminated unions may be sufficient.
• Document Your State Machine: Clearly document the different states, transitions, and actions of your state machine to improve maintainability and collaboration.

Advanced Techniques

Conditional Types

Conditional types can be combined with discriminated unions to create even more powerful and flexible state machines. For example, you can use conditional types to define different return types for a function based on the current state.

function getData(state: UIState): T | undefined {
  if (state.type === "success") {
    return state.data;
  }
  return undefined;
}

This function uses a simple `if` statement but could be made more robust using conditional types to ensure a specific type is always returned.

Utility Types

TypeScript's utility types, such as `Extract` and `Omit`, can be helpful when working with discriminated unions. `Extract` allows you to extract specific members from a union type based on a condition, while `Omit` allows you to remove properties from a type.

// Extract the "success" state from the UIState union
type SuccessState = Extract, { type: "success" }>;

// Omit the 'message' property from the Error interface
type ErrorWithoutMessage = Omit;

Real-World Examples Across Different Industries

The power of discriminated unions extends across various industries and application domains:

• E-commerce (Global): In a global e-commerce platform, order status can be represented with discriminated unions, handling states like "PaymentPending", "Processing", "Shipped", "InTransit", "Delivered", and "Cancelled". This ensures correct tracking and communication across different countries with varying shipping logistics.
• Financial Services (International Banking): Managing transaction states such as "PendingAuthorization", "Authorized", "Processing", "Completed", "Failed" is critical. Discriminated unions provide a robust way to handle these states, adhering to diverse international banking regulations.
• Healthcare (Remote Patient Monitoring): Representing patient health status using states like "Normal", "Warning", "Critical" enables timely intervention. In globally distributed healthcare systems, discriminated unions can ensure consistent data interpretation regardless of location.
• Logistics (Global Supply Chain): Tracking shipment status across international borders involves complex workflows. States like "CustomsClearance", "InTransit", "AtDistributionCenter", "Delivered" are perfectly suited for discriminated union implementation.
• Education (Online Learning Platforms): Managing course enrollment status with states such as "Enrolled", "InProgress", "Completed", "Dropped" can provide a streamlined learning experience, adaptable to different educational systems worldwide.

Conclusion

TypeScript discriminated unions provide a powerful and type-safe way to build state machines. By clearly defining the possible states and transitions, you can create more robust, maintainable, and understandable code. The combination of type safety, exhaustiveness checking, and enhanced code completion makes discriminated unions an invaluable tool for any TypeScript developer dealing with complex state management. Embrace discriminated unions in your next project and experience the benefits of type-safe state management firsthand. As we have shown with diverse examples from e-commerce to healthcare, and logistics to education, the principle of type-safe state management through discriminated unions is universally applicable.

Whether you're building a simple UI component or a complex enterprise application, discriminated unions can help you manage state more effectively and reduce the risk of runtime errors. So, dive in and explore the world of type-safe state machines with TypeScript!